#!/usr/bin/env python
# -*- coding=utf-8 -*-
"""
@author: xingwg
@license: (C) Copyright 2020-2025.
@contact: xingweiguo@chinasvt.com
@project: CenterNet
@file: convert_svt_to_coco.py
@time:  下午4:10
@desc:
"""
import os
import xmltodict
import json


xml_dir = "/media/xingwg/Run/Projects/detection/CenterNet/data/voc/VOCdevkit/VOC2007/Annotations"
list_file = "/media/xingwg/Run/Projects/detection/CenterNet/data/voc/VOCdevkit/VOC2007/ImageSets/Main/trainval.txt"
json_dir = "/media/xingwg/Run/Projects/detection/CenterNet/data/voc/VOCdevkit/VOC2007/Annotations_json"

if not os.path.isfile(list_file):
    exit(-1)

if not os.path.exists(xml_dir):
    exit(-1)

if not os.path.exists(json_dir):
    os.makedirs(json_dir)

attr = dict()
attr["categories"] = [{"supercategory": "none", "id": 1, "name": "car"},
                      {"supercategory": "none", "id": 2, "name": "pedestrian"},
                      {"supercategory": "none", "id": 3, "name": "rider"},
                      {"supercategory": "none", "id": 4, "name": "tricycle"}]

f = open(list_file, "r")
lines = f.readlines()
f.close()

phase = list_file.split("/")[-1].split(".")[0]
json_path = os.path.join(json_dir, phase + ".json")

file_names = list()
for line in lines:
    file_names.append(line.strip())

images = list()
annotations = list()
image_id = 0
for file_name in file_names:

    xml_path = os.path.join(xml_dir, file_name + ".xml")

    if not os.path.isfile(xml_path):
        continue

    doc = xmltodict.parse(open(xml_path).read())

    image_id += 1

    image = dict()
    image["file_name"] = str(doc["annotation"]["filename"])
    image["height"] = int(doc["annotation"]["size"]["height"])
    image["width"] = int(doc["annotation"]["size"]["width"])
    image['id'] = image_id

    images.append(image)

    obj_id = 1
    if 'object' in doc['annotation']:
        for obj in doc['annotation']['object']:
            for value in attr["categories"]:
                annotation = dict()
                if str(obj['name']) == value["name"]:
                    annotation["iscrowd"] = 0
                    annotation["image_id"] = image_id
                    x = int(obj["bndbox"]["xmin"])
                    y = int(obj["bndbox"]["ymin"])
                    w = int(obj["bndbox"]["xmax"]) - x
                    h = int(obj["bndbox"]["ymax"]) - y
                    annotation["bbox"] = [x, y, w, h]
                    annotation["area"] = float(w * h)
                    annotation["category_id"] = value["id"]
                    annotation["ignore"] = 0
                    annotation["id"] = obj_id
                    annotation["segmentation"] = []
                    obj_id += 1

                    annotations.append(annotation)

    attr["images"] = images
    attr["annotations"] = annotations
    attr["type"] = "instances"

    json_str = json.dumps(attr)
    with open("json_path", "w") as f:
        f.write(json_str)
